import platgo as pg
import numpy as np


class MOTSP(pg.Problem):

    def __init__(self, D: int = 30, M: int = 2, c: float = 0, C: np.ndarray = None):
        # 对于多目标的TSP问题，矩阵C应该是三维的，内层的二维矩阵应该是M个等大小的方针
        self.name = "MOTSP"
        self.borders = []
        self.D = D
        self.M = M
        self.type['multi'], self.type['many'], self.type['permutation'], self.type['large'] = [True] * 4
        if C is None:
            self.C = np.random.random((self.M, self.D, self.D))
            for i in range(1, M):
                self.C[i] = c * self.C[i-1] + (1 - c) * self.C[i]
        else:
            self.C = C
            self.M = C.shape[0]
            self.D = C.shape[-1]
            for i in range(self.M):
                self.C[i] = np.tril(self.C[i], k=-1) + np.triu(self.C[i].T, k=1)
        super().__init__()

    def cal_obj(self, pop: pg.Population) -> None:
        N, M = pop.decs.shape
        pop.objv = np.zeros((N, self.M))
        for i in range(self.M):
            for j in range(N):
                for k in range(self.D-1):
                    pop.objv[j][i] = pop.objv[j][i] + self.C[i][pop.decs[j][k]][pop.decs[j][k+1]]
                pop.objv[j][i] = pop.objv[j][i] + self.C[i][pop.decs[j][M-1]][pop.decs[j][0]]

    def get_optimal(self) -> np.ndarray:
        # generate a point for hypervolume calculation
        point = np.zeros((1, self.M)) + self.D
        return point


if __name__ == "__main__":
    C=np.array([[[0.0596188675796392,	0.0967300257808670,	0.659605252908307,	0.453797708726920,	0.173388613119006],
[0.681971904149063,	0.818148553859625,	0.518594942510538,	0.432391503783462,	0.390937802323736],
[0.0424311375007417,	0.817547092079286,	0.972974554763863,	0.825313795402046,	0.831379742839070],
[0.0714454646006424,	0.722439592366842,	0.648991492712356,	0.0834698148589140,	0.803364391602440],
[0.521649842464284,	0.149865442477967,	0.800330575352402,	0.133171007607162,	0.0604711791698936]],
[[0.399257770613576, 0.291984079961715,	0.106216344928664,	0.951630464777727,	0.422835615008808],
[0.526875830508296,	0.431651170248720,	0.372409740055537,	0.920332039836564,	0.547870901214845],
[0.416799467930787,	0.0154871256360190,	0.198118402542975,	0.0526769976807926,	0.942736984276934],
[0.656859890973707,	0.984063724379154,	0.489687638016024,	0.737858095516997,	0.417744104316662],
[0.627973359190104,	0.167168409914656,	0.339493413390758,	0.269119426398556,	0.983052466469856]]])

    a = MOTSP(D=5, C=C)
    print(a.C)
    decs = np.array([[1, 2, 3, 4, 5],
               [5, 4, 3, 2, 1],
               [2, 3, 4, 5, 1],
               [5, 2, 1, 4, 3],
               [1, 4, 5, 2, 3]])
    decs = decs - 1
    pop = pg.Population(decs=decs)
    a.cal_obj(pop=pop)
    print(pop.objv)
